package com.lfy.ai.kimi;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.utils.Constants;
import com.alibaba.dashscope.utils.JsonUtils;
import com.lfy.ai.AiChatStandard;
import com.lfy.ai.AiGroupManager;
import com.lfy.ai.AiGroupNorm;
import com.lfy.ai.Key;
import com.lfy.ai.kimi.sdk.Message;
import com.lfy.ai.kimi.sdk.MoonshotAiUtils;
import com.lfy.ai.weny.WenyGroupDialogue;
import com.lfy.domain.BuilderChatResult;
import com.lfy.modules.ai.cache.AiDispatchCache;
import com.lfy.modules.common.config.sys.DictionariesKey;
import com.lfy.modules.common.config.sys.SysDictionaries;
import com.lfy.modules.common.custom.constants.AIEnum;
import com.lfy.modules.common.domain.AiChatVO;
import com.lfy.modules.common.domain.BaseChatVO;
import com.lfy.websocket.cache.ChannelCache;
import io.netty.channel.Channel;
import lombok.extern.slf4j.Slf4j;

import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Import;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;

@Slf4j
@Component
@ConditionalOnProperty(name = "ai.kimi.enable", havingValue = "true", matchIfMissing = false)
@Import({KimiGroupDialogue.class})
public class KimiChat extends AiChatStandard {
    {
        log.info("Blog-web AI KIMI模型启动中...");
    }
    @Resource
    private AiDispatchCache aiDispatchCache;

    @Override
    protected boolean choice(Integer aiCode) {
        return AIEnum.AiModel.KimiAi.getCode().equals(aiCode);
    }

    @Override
    public void ask(String question, Integer modelType, String clientId, Integer groupId, String url) {
        AiGroupNorm group = AiGroupManager.getAiGroup(AIEnum.AiModel.KimiAi.getCode());
        if(Objects.isNull(group)){
            log.error("KIMI in ask method,group is not exist,modelId:{},groupId:{}",modelType,groupId);
            throw new IllegalArgumentException("group not found");

        }
        com.alibaba.dashscope.common.Message userMsg = com.alibaba.dashscope.common.Message.builder().role(Role.USER.getValue()).content(question).build();
        group.putMessage(clientId,groupId,userMsg);
        List<com.alibaba.dashscope.common.Message> messages = group.getMessages(clientId, groupId);
        //转换为KIMI用的message
        List<Message> messageList = messages.stream().map(item -> {
            Message message = new Message();

            message.setRole(item.getRole());

            message.setContent(item.getContent());

            return message;
        }).collect(Collectors.toList());

        MoonshotAiUtils.chat(AIEnum.KIMIType.getKIMIType(modelType), messageList, new kimiCallback() {
            @Override
            public void onResult(String result) {
                Channel channel = ChannelCache.getChannel(clientId);

                if (channel!= null && channel.isActive()) {
                    AiChatVO aiChatVO = new AiChatVO();
                    aiChatVO.setContent(result);
                    aiChatVO.setModelId(AIEnum.AiModel.KimiAi.getCode());
                    aiChatVO.setGroupId(groupId);
                    BaseChatVO<AiChatVO> chattingVO = BaseChatVO.getChattingVO(aiChatVO);
                    channel.writeAndFlush(BuilderChatResult.getAIResponseVO(chattingVO));
                }
            }

            @Override
            public void onError(String error) {
                log.error("KIMI交互错误 Exception occurred: {}", error);
                Channel channel = ChannelCache.getChannel(clientId);
                if (channel!= null && channel.isActive()) {
                    channel.writeAndFlush(BuilderChatResult.getAiExceptionResponseVO(error,groupId,AIEnum.AiModel.KimiAi.getCode()));
                }
            }

            @Override
            public void onCompletion() {
                Channel channel = ChannelCache.getChannel(clientId);
                if (channel!= null && channel.isActive()) {
                    channel.writeAndFlush(BuilderChatResult.getAiEndResponseVO(groupId,AIEnum.AiModel.KimiAi.getCode()));
                }
            }

            @Override
            public void onAllResult(String allResult) {
                try{
                    //保存系统回复消息
                    com.alibaba.dashscope.common.Message message = new com.alibaba.dashscope.common.Message();
                    message.setRole(Role.ASSISTANT.getValue());
                    message.setContent(allResult);
                    group.putMessage(clientId,groupId,message);

                    List<com.alibaba.dashscope.common.Message> list = group.getMessages(clientId, groupId);
                    if (CollectionUtils.isEmpty(list)) {
                        log.error("KIMI in ask method,list is empty,modelId:{},groupId:{}",modelType,groupId);
                        return;
                    }
                    tokenizerAfter( AIEnum.AiModel.KimiAi.getCode(),list, url);
                }catch (Exception e){
                    log.error("KIMI tokenizerAfter error:{},url:{},question:{},modelType:{}",e,url,allResult,modelType);
                }

            }
        });
    }

    @Override
    public Integer tokenizer(String question, Integer aiModelType) {
        Message message = new Message();
        message.setRole("user");
        message.setContent(question);
        String count = MoonshotAiUtils.estimateTokenCount(AIEnum.KIMIType.getKIMIType(aiModelType), Arrays.asList(message));

        Optional<Integer> integer = Optional.ofNullable(JSONUtil.parseObj(count)).map(e -> (JSONObject) e.get("data")).map(e -> (Integer) e.get("total_tokens"));

        log.info("KIMI 用户提问计算出来的token数量为:{}",integer.get());
        return integer.get();
    }
    private Integer getSysToken(List<com.alibaba.dashscope.common.Message> messages, Integer aiModelType){
        List<Message> result = messages.stream().map(item -> {
            Message message = new Message();
            message.setRole(item.getRole());
            message.setContent(item.getContent());
            return message;
        }).collect(Collectors.toList());
        String count = MoonshotAiUtils.estimateTokenCount(AIEnum.KIMIType.getKIMIType(aiModelType),result);
        Optional<Integer> integer = Optional.ofNullable(JSONUtil.parseObj(count))
                .map(e -> (JSONObject) e.get("data"))
                .map(e -> (Integer) e.get("total_tokens"));

        log.info("KIMI系统回答计算出来的token数量为:{}",integer.get());
        return integer.get();
    }
    @Override
    protected void tokenizerAfter(Integer aiModelType,List<com.alibaba.dashscope.common.Message> messages,  String url) {
        Integer tokenizer = getSysToken(messages, aiModelType);
        Long userId = aiDispatchCache.getUserId(url);
        aiDispatchCache.updateUserCapacity(AIEnum.AiModel.KimiAi.getCode(), aiModelType,userId,Long.valueOf(tokenizer));
        aiDispatchCache.updateUserCapacityCache(AIEnum.AiModel.KimiAi.getCode(), aiModelType,userId,Long.valueOf(tokenizer));
    }

}
